"""
@Description :   Gounding DINO Base 模型
@Author      :   tqychy 
@Time        :   2025/08/30 18:11:14
"""
import torch
from PIL import Image
from transformers import AutoModelForZeroShotObjectDetection, AutoProcessor


class GroundingDinoBase:
    def __init__(self, *args, model_path):
        self.cfg, self.logger = args
        self.model = AutoModelForZeroShotObjectDetection.from_pretrained(
            model_path,
            device_map="auto",
            local_files_only=True
        )
        self.processor = AutoProcessor.from_pretrained(model_path)

    def __call__(self, inputs):
        image_path, prompt = inputs
        image = Image.open(image_path).convert("RGB")
        inputs = self.processor(images=image, text=prompt, return_tensors="pt").to("cuda")
        with torch.no_grad():
            outputs = self.model(**inputs)
        results = self.processor.post_process_grounded_object_detection(
            outputs,
            inputs.input_ids,
            # box_threshold=0.4,
            text_threshold=0.3,
            target_sizes=[image.size[::-1]]
        )
        return results[0]

    @staticmethod
    def make_prompt(image_path, sentence, category, mode):
        if mode == "rec":
            prompt = sentence
        elif mode == "detect":
            prompt = category
        else:
            raise ValueError(f"Unknown task type! {mode}")
        return image_path, prompt

    @staticmethod
    def covert_formatted_bbox(bbox, image_shape):
        """
        将 [x1, y1, x2, y2] 格式的 bbox 转换成 [x_min, y_min, width, height] 格式的​
        """
        x1, y1, x2, y2 = bbox
        x_min = min(x1, x2)
        y_min = min(y1, y2)
        width = max(x1, x2) - x_min
        height = max(y1, y2) - y_min
        return [x_min, y_min, width, height]